/*
  Copyright 2006-2012 Patrik Jonsson, sunrise@familjenjonsson.org

  This file is part of Sunrise.

  Sunrise is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 3 of the License, or
  (at your option) any later version.

  Sunrise is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.
  
*/

#include "analytic_case.h"
#include "blackbody.h"
#include "blackbody_grain.h"
#include "misc.h"
#include "wd01_Brent_PAH_grain_model.h"

class irdc_factory {
public:
  typedef T_grid::T_grid_impl::T_data T_data;
  typedef refinement_accuracy_data<T_data> T_racc;
  typedef cell_tracker<T_data> T_cell_tracker;

  const T_densities rho0_, n_tol_; // central density and refinement tolerance
  const T_float rtol_, rsz_, xzr_, yzr_, rpow_; // radius tolerance and scale radius in z, and x/z, y/z scale ratios, radial power
  const T_densities rho_bg_; // background density
  const int maxl_, n_lambda_; // maxlevel and n_lambda
  const int n_threads_;
  
  irdc_factory (const T_densities& rho0, const T_densities& n_tol, T_float rtol,
		T_float rsz, T_float xzr, T_float yzr, T_float rpow, 
		T_densities rhobg, int ml, int nl, int nt):
    rho0_(rho0), n_tol_(n_tol), rtol_(rtol), rsz_(rsz), xzr_(xzr), yzr_(yzr), 
    rpow_(rpow),
    rho_bg_(rhobg), maxl_(ml), n_lambda_(nl), n_threads_(nt) {};

  T_densities calc_rho(const vec3d& p) {
    T_densities rho;
    resize_like(rho, rho0_);
    rho = rho_bg_;
    // elliptical radius
    const T_float r2 = p[2]*p[2] + p[0]*p[0]/(xzr_*xzr_) + 
      p[1]*p[1]/(yzr_*yzr_);
    rho += rho0_/(1+pow(r2/(rsz_*rsz_),rpow_/2));
    return rho;
  };

  bool refine_cell_p (const T_cell_tracker& c) {
    const int level = c.code().level();
    if (level>maxl_) 
      return false;
    //if (level<=4) return true; // a minimum amount of "prerefinement"
    // refine until cell column density < n_tol
    
    // simone uses maxlevel=10+log(r/rmin)
    const vec3d min = c.getmin();
    const vec3d max = c.getmax();
    const T_float sz=sqrt(dot(c.getsize(),c.getsize()));
    const T_float r=sqrt(dot(c.getcenter(), c.getcenter()));

    T_densities rhoav
      (-log((exp(-sz*calc_rho(c.getcenter()))+
	     exp(-sz*calc_rho(min))+
	     exp(-sz*calc_rho(vec3d(min[0],min[1],max[2])))+
	     exp(-sz*calc_rho(vec3d(min[0],max[1],min[2])))+
	     exp(-sz*calc_rho(vec3d(min[0],max[1],max[2])))+
	     exp(-sz*calc_rho(vec3d(max[0],min[1],min[2])))+
	     exp(-sz*calc_rho(vec3d(max[0],min[1],max[2])))+
	     exp(-sz*calc_rho(vec3d(max[0],max[1],min[2])))+
	     exp(-sz*calc_rho(max)))/9)/sz) ;    

    if (any(rhoav>2*rho_bg_) && (sz/r>rtol_))
      return true;


    // refine cell if column density exceeds n_tol
    return condition_all (rhoav*rhoav*dot(c.getsize(),c.getsize()) > 
			  n_tol_*n_tol_);
  };

  bool unify_cell_p (const T_cell_tracker& c, const T_racc& racc) {
    return false;
  };

  T_data get_data (const T_cell_tracker& c) {
    const vec3d min = c.getmin();
    const vec3d max = c.getmax();
    /*
    T_densities rhoav
      ((calc_rho(c.getcenter())+
	calc_rho(min)+
	calc_rho(vec3d(min[0],min[1],max[2]))+
	calc_rho(vec3d(min[0],max[1],min[2]))+
	calc_rho(vec3d(min[0],max[1],max[2]))+
	calc_rho(vec3d(max[0],min[1],min[2]))+
	calc_rho(vec3d(max[0],min[1],max[2]))+
	calc_rho(vec3d(max[0],max[1],min[2]))+
	calc_rho(max))/9);
    */
    const T_float sz=sqrt(dot(c.getsize(),c.getsize()));
    T_densities rhoav
      (-log((exp(-sz*calc_rho(c.getcenter()))+
	     exp(-sz*calc_rho(min))+
	     exp(-sz*calc_rho(vec3d(min[0],min[1],max[2])))+
	     exp(-sz*calc_rho(vec3d(min[0],max[1],min[2])))+
	     exp(-sz*calc_rho(vec3d(min[0],max[1],max[2])))+
	     exp(-sz*calc_rho(vec3d(max[0],min[1],min[2])))+
	     exp(-sz*calc_rho(vec3d(max[0],min[1],max[2])))+
	     exp(-sz*calc_rho(vec3d(max[0],max[1],min[2])))+
	     exp(-sz*calc_rho(max)))/9)/sz) ;
    //const T_densities rhoav = calc_rho(c.getcenter());
    return T_data(T_data::T_emitter(), 
		  T_data::T_absorber(rhoav, vec3d(0,0,0), T_lambda(n_lambda_)));
  }
    
  int n_threads () {return n_threads_;};
};

T_float unit_opacity_rho;


void setup (po::variables_map& opt)
{
  T_unit_map units;
  units ["length"] = "pc";
  units ["mass"] = "Msun";
  units ["luminosity"] = "W";
  units ["wavelength"] = "m";
  units ["L_lambda"] = "W/m";

  const int maxlevel= opt["maxlevel"].as<int>();
  const T_float grid_extent = opt["gridsize"].as<T_float>();
  const T_float cameradist = 1e5;
  const int npix = opt["npix"].as<int>();
  const int n_threads = opt["n_threads"].as<int>();
  const string grain_model = opt["grain_model"].as<string>();

  const T_float rsz=opt["rsz"].as<T_float>();
  const T_float xzr=opt["xzr"].as<T_float>();
  const T_float yzr=opt["yzr"].as<T_float>();
  const T_float rpow=opt["rpow"].as<T_float>();

  const T_float rho0 = opt["rho0"].as<T_float>();
  const T_float rhobg = opt["rho_bg"].as<T_float>();
  const T_float tau_tol = opt["tau_tol"].as<T_float>();
  const T_float r_tol = opt["r_tol"].as<T_float>();

  const T_float U=opt["intensity"].as<T_float>();

  Preferences p;
  p.setValue("grain_data_directory", 
		 word_expand(opt["grain_data_dir"].as<string>())[0]);
  p.setValue("grain_model", grain_model);
  p.setValue("wd01_parameter_set", opt["wd01_parameter_set"].as<string>());
  p.setValue("template_pah_fraction", 0.5);
  p.setValue("use_dl07_opacities", true);
  p.setValue("n_threads", n_threads);
  p.setValue("use_grain_temp_lookup", true);
  p.setValue("use_cuda", false);

  // opacity for the WD01 dust = 2.41e-6 kpc^2/Msun (=1.025e11 au^2/Msun)
  // which implies tau = 3.54e14 au^3/Msun * rho0

  array_1 lambda;
  lambda.reference(logspace(opt["lambda_min"].as<T_float>(),
			    opt["lambda_max"].as<T_float>(),
			    opt["n_lambda"].as<int>()));

  // load grain model
  T_dust_model::T_scatterer_vector sv;
  sv.push_back(load_grain_model<polychromatic_scatterer_policy, mcrx_rng_policy>(p,units));
  T_dust_model model (sv);
  model.set_wavelength(lambda); // this also resamples the grain_model objects

  std::vector<T_float> lambdas(lambda.begin(),lambda.end());
  const int lambda550  = lower_bound(lambdas.begin(), lambdas.end(),
				     0.551e-6) - lambdas.begin()-1;
  
  // Convert optical depth to density
  unit_opacity_rho = 1.0/sv[0]->opacity()(lambda550);
  T_densities rho_0(1); rho_0=rho0;
  T_densities rho_bg(1); rho_bg=rhobg;

  T_densities n_tol(1); n_tol = tau_tol*unit_opacity_rho;
  cout << "Unit opacity density: " << unit_opacity_rho <<  ", n_tol: " << n_tol(0) << endl;
  //ONLY VALID FOR RPOW=2
  cout << "Estimated V-band optical depth through center " << 2*rho0*rsz*atan(grid_extent/rsz)/unit_opacity_rho << endl;

  // build grid
  cout << "Building grid" << endl;
  irdc_factory factory (rho_0,n_tol, r_tol, rsz, xzr, yzr, rpow,
			rho_bg, maxlevel, 
			lambda.size(), n_threads);

  // calculate luminosity necessary to get 1U intensity from ext field  
  const T_float r_em=1.5*sqrt(3)*grid_extent;
  const T_float L_em=4*constants::pi*pow(units::convert(units.get("length"), "m")*r_em,2)*constants::c0/3*U*8.65e-14;
  cout << "External illumination at radius " << r_em << ", luminosity " << L_em << endl;
  blackbody bb(5800,L_em);

  boost::shared_ptr<T_emission> em;
  //em.reset(new pointsource_emission<polychromatic_policy, T_rng_policy> (vec3d (.01, .01, -.01), bb.emission(lambda)));
  em.reset(new external_illumination<polychromatic_policy, T_rng_policy> (r_em, sqrt(3)*grid_extent, bb.emission(lambda)));
  //em.reset(new divergent_floodlight_emission<polychromatic_policy, T_rng_policy> (vec3d(0,0,-500),vec3d(0,0,1),100,0.5, bb.emission(lambda)));

  std::vector<std::pair<T_float, T_float> > cam_pos;
  cam_pos.push_back(std:: make_pair (0.*constants::pi/180, 0.) );
  cam_pos.push_back(std:: make_pair (90*constants::pi/180, 0.) );
  cam_pos.push_back(std:: make_pair (90*constants::pi/180, 90*constants::pi/180) );
  boost::shared_ptr<T_emergence> e
    (new T_emergence(cam_pos, cameradist, 2*grid_extent, npix));

  run_case(opt, 
	   factory,
	   units,
	   lambda,
	   vec3d(-grid_extent,-grid_extent,-grid_extent),
	   vec3d(grid_extent,grid_extent,grid_extent),
	   *em,
	   *e,
	   model); 
}

void pre_shoot(T_xfer& x)
{
  T_densities n_col(x.integrate_column_density(vec3d(0,0,0),vec3d(0,0,1)));
  std::cout << "Actual V-band optical depth through center: " << 2*n_col(0)/unit_opacity_rho << std::endl;
}

std::string description()
{
  return "IRDC model.";
}

void add_custom_options(po::options_description& desc)
{
  desc.add_options()
    ("gridsize", po::value<T_float>()->default_value(10),
     "Grid size (pc)")
    ("intensity", po::value<T_float>()->default_value(1),
     "Intensity of external radiation (U)")
    ("rsz", po::value<T_float>()->default_value(0.08),
     "Cloud scale radius in z direction (pc)")
    ("xzr", po::value<T_float>()->default_value(1),
     "Ratio of x- and z- cloud scale radii")
    ("yzr", po::value<T_float>()->default_value(1),
     "Ratio of y- and z- cloud scale radii")
    ("rpow", po::value<T_float>()->default_value(2.0),
     "Exponent of density radial dependence")
    ("rho0", po::value<T_float>()->default_value(1e3),
     "Central cloud dust density (Msun/pc^3)")
    ("rho_bg", po::value<T_float>()->default_value(1e-3),
     "Background dust density (Msun/pc^3)")
    ("r_tol", po::value<T_float>()->default_value(.25),
     "cell fractional size/radius tolerance")
    ("grain_model", po::value<string>()->default_value("wd01_Brent_PAH"), "dust grain model (use \"blackbody\" for blackbody grains)")
    ;
}

void print_custom_options(const po::variables_map& opt)
{
  std::cout
       << "gridsize = " << opt["gridsize"].as<T_float>() << '\n'
       << "intensity = " << opt["intensity"].as<T_float>() << '\n'
       << "rsz = " << opt["rsz"].as<T_float>() << '\n'
       << "xzr = " << opt["xzr"].as<T_float>() << '\n'
       << "yzr = " << opt["yzr"].as<T_float>() << '\n'
       << "rpow = " << opt["rpow"].as<T_float>() << '\n'
       << "rho0 = " << opt["rho0"].as<T_float>() << '\n'
       << "rho_bg = " << opt["rho_bg"].as<T_float>() << '\n'
       << "r_tol = " << opt["r_tol"].as<T_float>() << '\n'
       << "grain_model = " << opt["grain_model"].as<std::string>() << '\n'
    ;
}
